package com.alibaba.datax.plugin.writer.greenplumwriter;

import org.apache.commons.lang3.StringUtils;
import org.postgresql.copy.CopyManager;
import org.postgresql.core.BaseConnection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.sql.DataSource;
import java.io.*;
import java.sql.Connection;
import java.sql.SQLException;

/**
 * @Description: CopyIn工具类
 * @Author: chenweifeng
 * @Date: 2022年08月11日 下午4:26
 **/
public class CopyInDataUtils {

    private static Logger logger = LoggerFactory.getLogger(CopyInDataUtils.class);


    /**
     * copy in导出数据
     *
     * @param dataSource 数据源
     * @param querySql   查询sql
     * @param delimiter  列分隔符 例如 ,
     * @param filePath   导出文件路径
     * @param encode     编码 例如 UTF-8
     * @param header     是否需要表头
     * @return
     */
    public static long exportData(DataSource dataSource,
                                  String querySql,
                                  String delimiter,
                                  String filePath,
                                  String encode,
                                  boolean header) {
        Connection con = null;
        OutputStream out = null;
        Writer writer = null;
        try {
            con = dataSource.getConnection();
            CopyManager cm = new CopyManager(con.unwrap(BaseConnection.class));
            StringBuffer sb = new StringBuffer();
            sb.append("copy (");
            sb.append(querySql);
            sb.append(" ) TO STDOUT ");
            sb.append("WITH DELIMITER '");
            sb.append(delimiter);
            sb.append("'");
            if (header) {
                sb.append(" HEADER ");
            }
            String copySql = sb.toString();
            logger.info("exportData data begin ,  sql  is {}", copySql);

            long startTime = System.currentTimeMillis();
            File file = new File(filePath);
            if (!file.getParentFile().exists()) {
                file.getParentFile().mkdirs();
            }
            out = new FileOutputStream(file);
            long handledRowCount = 0;
            if (StringUtils.isNotEmpty(encode)) {
                writer = new OutputStreamWriter(out, encode);
                handledRowCount = cm.copyOut(copySql, writer);
            } else {
                handledRowCount = cm.copyOut(copySql, out);
            }
            long elapsedTime = System.currentTimeMillis() - startTime;

            logger.info("exportData data end, sql  is {}, elapsed time = {}", copySql, elapsedTime);
            return handledRowCount;
        } catch (Exception e) {
            logger.error(e.getMessage(), e);
            return 0L;
        } finally {
            if (out != null) {
                try {
                    out.close();
                } catch (IOException e) {
                    logger.error(e.getMessage(), e);
                }
            }

            if (writer != null) {
                try {
                    writer.close();
                } catch (IOException e) {
                    logger.error(e.getMessage(), e);
                }
            }

            if (con != null) {
                try {
                    con.close();
                } catch (SQLException e) {
                    logger.error(e.getMessage(), e);
                }
            }
        }
    }

    /**
     * copy in导入数据
     *
     * @param dataSource 导入的数据源
     * @param table      表名
     * @param delimiter  分隔符
     * @param filePath   导入数据的文件路径
     * @return
     */
    public static long importData(DataSource dataSource,
                                  String table,
                                  String delimiter,
                                  String filePath) {
        Connection con = null;
        InputStream in = null;
        try {
            logger.info("import data begin");
            con = dataSource.getConnection();
            CopyManager cm = new CopyManager(con.unwrap(BaseConnection.class));
            StringBuffer sb = new StringBuffer();
            sb.append("copy ");
            sb.append(table);
            sb.append(" from STDIN  ");
            sb.append("WITH DELIMITER '");
            sb.append(delimiter);
            sb.append("'");
            String copySql = sb.toString();
            logger.info("import data begin,  sql  is {}", copySql);
            long startTime = System.currentTimeMillis();
            in = new FileInputStream(filePath);
            long handledRowCount = cm.copyIn(copySql, in);
            long elapsedTime = System.currentTimeMillis() - startTime;
            logger.info("import data end,  sql  is {}, elapsed time = {}", copySql, elapsedTime);
            return handledRowCount;
        } catch (Exception e) {
            logger.error(e.getMessage(), e);
            return 0L;
        } finally {
            if (in != null) {
                try {
                    in.close();
                } catch (IOException e) {
                    logger.error(e.getMessage(), e);
                }
            }

            if (con != null) {
                try {
                    con.close();
                } catch (SQLException e) {
                    logger.error(e.getMessage(), e);
                }
            }
        }
    }
}
 
 